import cv2
import numpy as np
import os
import json

# 宏定义：棋盘格行数和列数
ROWS = 6
COLS = 9
# 宏定义：棋盘格的尺寸
SQUARE_SIZE = 20

JSON_FILE = True  # 是否将标定结果写入到 JSON 文件中
IMG_PATH = "/home/shuai/pyorbbecsdk-1.3.2/cam_robot/img"
JSON_PATH = "/home/shuai/pyorbbecsdk-1.3.2/cam_robot/json"


class CamCalibration:
    def __init__(self):
        self.chessboard_size = (COLS, ROWS)
        self.square_size = SQUARE_SIZE
        self.obj_points = []
        self.img_points = []
        self.objp = np.zeros((ROWS * COLS, 3), np.float32)
        self.objp[:, :2] = np.mgrid[0:COLS, 0:ROWS].T.reshape(-1, 2)
        self.objp = SQUARE_SIZE * self.objp

    def calibrate_camera(self, image_dir):
        image_files = [
            f for f in os.listdir(image_dir) if f.endswith(".jpg") or f.endswith(".png")
        ]
        for fname in image_files:
            img = cv2.imread(os.path.join(image_dir, fname))
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            ret, corners = cv2.findChessboardCorners(gray, self.chessboard_size, None)
            if ret:
                self.obj_points.append(self.objp)
                corners2 = cv2.cornerSubPix(
                    gray,
                    corners,
                    (5, 5),
                    (-1, -1),
                    (cv2.TERM_CRITERIA_MAX_ITER | cv2.TERM_CRITERIA_EPS, 30, 0.001),
                )
                self.img_points.append(corners2 if corners2 is not None else corners)
                cv2.drawChessboardCorners(img, self.chessboard_size, corners, ret)
                cv2.waitKey(1)
        _, mtx, dist, _, _ = cv2.calibrateCamera(
            self.obj_points, self.img_points, gray.shape[::-1], None, None
        )
        self.camera_intrinsic = {"mtx": mtx.tolist(), "dist": dist.tolist()}
        return self.camera_intrinsic


if __name__ == "__main__":
    calib = CamCalibration()
    intrinsic = calib.calibrate_camera(IMG_PATH)
    print(intrinsic)

    if JSON_FILE:
        if not os.path.exists("json"):
            os.makedirs("json")
            with open(os.path.join(JSON_PATH, "intrinsic_params.json"), "w") as f:
                json.dump(intrinsic, f, indent=4)
        print("Camera intrinsic parameters saved to 'intrinsic_params.json'")
